import os
import time
import logging
import subprocess
import sys

try:
    import torch

    # 检测CUDA是否可用，用于判断是否可以使用GPU加速
    CUDA_AVAILABLE = torch.cuda.is_available()
except ImportError:
    # 如果没有安装torch库，则默认不使用CUDA
    CUDA_AVAILABLE = False
# 导入LangChain相关模块用于构建聊天系统
from langchain_classic.chains import ConversationChain
from langchain_classic.memory import ConversationBufferWindowMemory
from langchain_core.prompts import PromptTemplate  # 提示模板                  # 提示模板
# 使用LlamaCpp加载本地GGUF格式的大语言模型
from langchain_community.llms import LlamaCpp
# 从HuggingFace下载模型文件
from huggingface_hub import hf_hub_download

# 导入Flask用于创建简单的MCP服务
try:
    from flask import Flask, request, jsonify

    FLASK_AVAILABLE = True
except ImportError:
    FLASK_AVAILABLE = False
    Flask = None
    request = None
    jsonify = None

"""
兼容CPU和CUDA的本地Qwen大语言模型聊天程序
支持自动检测系统环境并在GPU或CPU上运行模型
支持MCP服务接入
"""

# 配置网络代理（如果需要）
os.environ["http_proxy"] = "http://127.0.0.1:7890"
os.environ["https_proxy"] = "http://127.0.0.1:7890"

# 配置日志系统，用于输出程序运行信息
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# 配置参数 - 针对GPU/CPU优化
MODEL_REPO = "Qwen/Qwen1.5-1.8B-Chat-GGUF"  # 使用1.8B量化版本
MODEL_FILE = "qwen1_5-1_8b-chat-q4_0.gguf"  # 4位量化模型
MODEL_PATH = "E:\\LLMModel\\"  # 模型存储目录（Windows）
# MODEL_PATH = "/home/wsm/.lmstudio/models/Qwen1.5"  # 模型存储目录（Linux）
HISTORY_SIZE = 3  # 减少历史记录以降低内存占用，仅保留最近3轮对话
MAX_TOKENS = 256  # 减少生成长度加速响应，最大生成256个token
N_THREADS = os.cpu_count()  # 动态获取CPU核心数
CONTEXT_SIZE = 2048  # 上下文长度，模型能处理的最大上下文token数

# MCP配置参数
MCP_HOST = "localhost"  # MCP服务主机地址
MCP_PORT = 3000  # MCP服务端口
MCP_ENABLED = True  # 是否启用MCP服务
MCP_EXTERNAL = False  # 是否使用外部MCP服务（而非内置Flask服务）
MCP_LOCAL_PATH = "/home/wsm/codes/mcp_server/mcp_server.py"  # 本地MCP服务器路径


# 1. 下载并加载量化模型 (GGUF格式)
def load_cuda_optimized_model():
    """
    加载经过CUDA优化的本地大语言模型
    自动检测系统环境并在GPU或CPU上运行模型
    """
    logger.info("⏳ 正在下载并加载CUDA优化模型...")

    # 确保模型目录存在，如果不存在则创建
    os.makedirs(MODEL_PATH, exist_ok=True)
    local_model_path = os.path.join(MODEL_PATH, "Qwen\\Qwen1.5-1.8B-Chat-GGUF", MODEL_FILE)

    # 下载模型文件 (如果本地不存在)
    if not os.path.exists(local_model_path):
        logger.info(f"  下载模型: {MODEL_FILE}...")
        try:
            # 从HuggingFace Hub下载模型文件
            model_path = hf_hub_download(
                repo_id=MODEL_REPO,  # 模型仓库ID
                filename=MODEL_FILE,  # 模型文件名
                cache_dir=MODEL_PATH,  # 缓存目录
                resume_download=True  # 支持断点续传
            )
            logger.info(f"  模型下载完成: {model_path}")
        except Exception as e:
            logger.error(f"  模型下载失败: {str(e)}")
            raise RuntimeError(f"模型下载失败: {str(e)}")
    else:
        # 如果本地已存在模型文件，则直接使用
        logger.info(f"  使用本地模型: {local_model_path}")
        model_path = local_model_path

    # 检查模型文件大小并输出日志
    file_size = os.path.getsize(model_path) / (1024 * 1024)  # MB
    logger.info(f"  模型文件大小: {file_size:.2f} MB")

    # 配置CUDA优化参数
    try:
        logger.info("  尝试加载模型...")
        # 根据CUDA可用性设置n_gpu_layers参数
        # -1表示使用所有可用的GPU层，0表示不使用GPU（纯CPU模式）
        n_gpu_layers = -1 if CUDA_AVAILABLE else 0

        # 根据CUDA可用性输出相关信息
        if CUDA_AVAILABLE:
            logger.info("✅ 检测到CUDA，将在GPU上运行模型")
            logger.info(f"   CUDA版本: {torch.version.cuda}")
            logger.info(f"   GPU型号: {torch.cuda.get_device_name()}")
        else:
            logger.info("⚠️ 未检测到CUDA，将在CPU上运行模型")

        # 创建LlamaCpp实例并配置相关参数
        llm = LlamaCpp(
            model_path=model_path,  # 模型文件路径
            n_ctx=CONTEXT_SIZE,  # 上下文长度
            max_tokens=MAX_TOKENS,  # 最大生成token数
            temperature=0.7,  # 控制生成文本的随机性，值越高越随机
            n_gpu_layers=n_gpu_layers,  # 根据CUDA可用性自动设置GPU层
            n_threads=N_THREADS,  # CPU线程数
            n_batch=512,  # 批处理大小，提高处理效率
            verbose=False,  # 是否输出详细日志
            use_mlock=True,  # 锁定模型在RAM中以提高性能
            seed=42,  # 随机种子，保证结果可复现
            f16_kv=True  # 使用半精度键值缓存，节省内存
        )
        logger.info("✅ 模型加载成功")
        return llm
    except Exception as e:
        logger.error(f"❌ 模型加载失败: {str(e)}")
        logger.error("请检查模型文件是否存在且格式正确")
        raise RuntimeError(f"模型加载失败: {str(e)}")


# 2. 创建轻量级聊天系统
def create_lightweight_chat_system(llm):
    """
    创建轻量级聊天系统
    使用ConversationChain和ConversationBufferWindowMemory管理对话历史
    """
    # 定义聊天提示模板，包含历史对话和当前输入
    template = """
    你是一个乐于助人的AI助手。请根据最近的对话历史回答用户问题。

    历史对话:
    {history}

    用户: {input}
    助手: 
    """

    # 创建PromptTemplate实例
    prompt = PromptTemplate(
        input_variables=["history", "input"],  # 模板中使用的变量
        template=template
    )

    # 创建带有限历史记忆的对话链
    # ConversationBufferWindowMemory只保留最近k轮对话，避免内存占用过大
    memory = ConversationBufferWindowMemory(
        k=HISTORY_SIZE,  # 仅保留最近3轮对话
        human_prefix="用户",  # 用户发言前缀
        ai_prefix="助手"  # AI助手发言前缀
    )

    # 返回配置好的对话链
    return ConversationChain(
        llm=llm,  # 使用的大语言模型
        prompt=prompt,  # 使用的提示模板
        memory=memory,  # 使用的对话记忆
        verbose=False  # 是否输出详细日志
    )


# 3. 初始化MCP服务
def initialize_mcp_service(chat_chain):
    """
    初始化并启动MCP服务
    为模型提供标准的HTTP接口服务
    根据配置决定使用内置Flask服务还是连接外部MCP服务
    """
    if not MCP_ENABLED:
        logger.info("MCP服务未启用")
        return None

    # 如果配置为使用外部MCP服务，则不需要启动内置服务
    if MCP_EXTERNAL:
        logger.info("使用外部MCP服务，无需启动内置服务")
        return "external"

    # 如果配置为使用本地MCP服务
    if os.path.exists(MCP_LOCAL_PATH):
        logger.info("使用本地MCP服务")
        try:
            # 在后台启动本地MCP服务
            mcp_process = subprocess.Popen([
                sys.executable, MCP_LOCAL_PATH
            ], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
            logger.info(f"✅ 本地MCP服务已在后台启动，PID: {mcp_process.pid}")
            return mcp_process
        except Exception as e:
            logger.error(f"❌ 本地MCP服务启动失败: {str(e)}")
            return None

    # 使用内置Flask服务
    if not FLASK_AVAILABLE:
        logger.warning("❌ Flask未安装，无法启动MCP服务")
        logger.warning("请运行: pip install flask")
        return None

    try:
        logger.info(f"正在初始化内置MCP服务: {MCP_HOST}:{MCP_PORT}")

        # 创建Flask应用实例
        app = Flask(__name__)

        # 禁用Flask的日志输出
        log = logging.getLogger('werkzeug')
        log.setLevel(logging.ERROR)

        # 定义聊天端点
        @app.route('/chat', methods=['POST'])
        def chat_endpoint():
            try:
                # 获取请求数据
                data = request.get_json()
                user_input = data.get("message", "")
                stream = data.get("stream", False)  # 是否流式输出

                if not user_input:
                    return jsonify({"error": "消息内容不能为空"}), 400

                # 使用聊天链处理用户输入
                result = chat_chain.invoke({"input": user_input})
                response = result["response"]

                return jsonify({
                    "response": response,
                    "status": "success"
                })
            except Exception as e:
                logger.error(f"处理聊天请求时出错: {str(e)}")
                return jsonify({"error": str(e)}), 500

        # 定义流式聊天端点
        @app.route('/chat/stream', methods=['POST'])
        def chat_stream_endpoint():
            try:
                # 获取请求数据
                data = request.get_json()
                user_input = data.get("message", "")

                if not user_input:
                    return jsonify({"error": "消息内容不能为空"}), 400

                def generate():
                    try:
                        # 流式生成响应
                        full_response = ""
                        for chunk in chat_chain.stream({"input": user_input}):
                            response = chunk.get("response", "")
                            # 只输出新增内容
                            new_content = response[len(full_response):]
                            if new_content:
                                yield new_content
                            full_response = response
                    except Exception as e:
                        logger.error(f"流式生成过程中出错: {str(e)}")
                        yield f"\n[错误: {str(e)}]"

                return app.response_class(generate(), mimetype='text/plain')
            except Exception as e:
                logger.error(f"处理流式聊天请求时出错: {str(e)}")
                return jsonify({"error": str(e)}), 500

        # 定义健康检查端点
        @app.route('/health', methods=['GET'])
        def health_endpoint():
            return jsonify({
                "status": "healthy",
                "model": "Qwen1.5-1.8B-Chat",
                "device": "GPU" if CUDA_AVAILABLE else "CPU",
                "cuda_available": CUDA_AVAILABLE
            })

        # 定义模型信息端点
        @app.route('/model/info', methods=['GET'])
        def model_info_endpoint():
            return jsonify({
                "model_name": "Qwen1.5-1.8B-Chat",
                "model_file": MODEL_FILE,
                "context_size": CONTEXT_SIZE,
                "max_tokens": MAX_TOKENS,
                "threads": N_THREADS,
                "device": "GPU" if CUDA_AVAILABLE else "CPU",
                "cuda_available": CUDA_AVAILABLE,
                "cuda_version": torch.version.cuda if CUDA_AVAILABLE else None,
                "gpu_name": torch.cuda.get_device_name() if CUDA_AVAILABLE else None
            })

        logger.info("✅ 内置MCP服务初始化完成")
        return app
    except Exception as e:
        logger.error(f"❌ MCP服务初始化失败: {str(e)}")
        return None


# 4. 流式响应生成器
def generate_response_stream(chain, user_input):
    """
    流式生成响应
    逐段生成并返回模型输出，提供更好的用户体验
    """
    full_response = ""
    # 使用chain.stream方法流式生成响应
    for chunk in chain.stream({"input": user_input}):
        response = chunk.get("response", "")
        # 只输出新增内容，避免重复输出
        new_content = response[len(full_response):]
        if new_content:
            yield new_content
        full_response = response


# 5. 主聊天函数 (CUDA优化版)
def cuda_chat_main():
    """
    主聊天函数
    控制整个聊天程序的流程
    """
    mcp_process = None
    try:
        # 加载模型并计算加载时间
        start_time = time.time()
        qwen_llm = load_cuda_optimized_model()
        load_time = time.time() - start_time
        logger.info(f"✅ 模型加载完成 (耗时: {load_time:.1f}秒)")

        # 创建聊天系统
        chat_chain = create_lightweight_chat_system(qwen_llm)
        logger.info(f"💬 聊天系统已就绪 (历史记忆: {HISTORY_SIZE}轮)")

        # 初始化MCP服务
        mcp_service = initialize_mcp_service(chat_chain)
        mcp_server = None
        if mcp_service and MCP_ENABLED and not MCP_EXTERNAL:
            # 检查是否为本地MCP进程
            if isinstance(mcp_service, subprocess.Popen):
                mcp_process = mcp_service
                logger.info(f"✅ 本地MCP服务已启动 PID: {mcp_process.pid}")
            else:
                # 在后台启动MCP服务
                import threading
                def run_mcp_server():
                    mcp_service.run(host=MCP_HOST, port=MCP_PORT, threaded=False)

                mcp_server = threading.Thread(target=run_mcp_server, daemon=True)
                mcp_server.start()
                logger.info(f"✅ 内置MCP服务已在 {MCP_HOST}:{MCP_PORT} 启动")
        elif MCP_EXTERNAL and MCP_ENABLED:
            logger.info("🔄 配置为使用外部MCP服务")

        # 根据设备类型显示不同的欢迎信息
        device_info = "GPU" if CUDA_AVAILABLE else "CPU"
        print(f"\n助手: 你好！我是运行在{device_info}上的Qwen助手，有什么可以帮你的吗？")
        if MCP_ENABLED and not MCP_EXTERNAL:
            if mcp_process:
                print("📡 本地MCP服务已启动")
            else:
                print(f"📡 内置MCP服务地址: {MCP_HOST}:{MCP_PORT}")
                print("   可用端点:")
                print("   - POST /chat         (发送消息)")
                print("   - POST /chat/stream  (流式发送消息)")
                print("   - GET /health        (健康检查)")
                print("   - GET /model/info    (模型信息)")
        elif MCP_ENABLED and MCP_EXTERNAL:
            print("🔄 已配置为使用外部MCP服务")

        # 聊天循环
        while True:
            try:
                # 获取用户输入
                user_input = input("\n你: ").strip()

                # 处理特殊命令 - 退出程序
                if user_input.lower() in ["/exit", "退出", "exit"]:
                    print("\n助手: 再见！感谢你的聊天~")
                    break

                # 处理特殊命令 - 清空对话历史
                if user_input == "/clear":
                    chat_chain.memory.clear()
                    print("助手: 记忆已清空，让我们重新开始吧！")
                    continue

                # 处理特殊命令 - 显示系统信息
                if user_input == "/info":
                    print("\n[系统信息]")
                    print(f"模型: {MODEL_FILE}")
                    print(f"线程: {N_THREADS}")
                    print(f"上下文: {CONTEXT_SIZE} tokens")
                    print(f"最大生成: {MAX_TOKENS} tokens")
                    print(f"设备: {'GPU (CUDA)' if CUDA_AVAILABLE else 'CPU'}")
                    if CUDA_AVAILABLE:
                        print(f"CUDA版本: {torch.version.cuda}")
                        print(f"GPU型号: {torch.cuda.get_device_name()}")
                    if MCP_ENABLED and not MCP_EXTERNAL:
                        if mcp_process:
                            print(f"MCP服务: 本地服务已启用 (PID: {mcp_process.pid})")
                        else:
                            print(f"MCP服务: 内置服务已启用 ({MCP_HOST}:{MCP_PORT})")
                    elif MCP_ENABLED and MCP_EXTERNAL:
                        print("MCP服务: 使用外部服务")
                    else:
                        print("MCP服务: 未启用")
                    continue

                # 生成并输出响应
                start_time = time.time()
                print("助手: ", end="", flush=True)

                # 流式输出响应内容
                full_response = ""
                for chunk in generate_response_stream(chat_chain, user_input):
                    print(chunk, end="", flush=True)
                    full_response += chunk

                # 计算并显示响应时间
                response_time = time.time() - start_time
                print(f"\n[响应时间: {response_time:.1f}秒]")

            except KeyboardInterrupt:
                # 处理Ctrl+C中断
                print("\n助手: 会话已中断")
                break
            except Exception as e:
                # 处理聊天过程中的其他异常
                logger.error(f"聊天过程中出错: {str(e)}")
                print("助手: 抱歉，处理这个问题时出错了，请换个问题试试？")

    except Exception as e:
        # 处理程序初始化过程中的异常
        logger.error(f"程序初始化失败: {str(e)}")
        print(f"❌ 程序启动失败: {str(e)}")
        print("请检查模型文件是否完整且格式正确")
    finally:
        # 清理资源，关闭MCP进程
        if mcp_process and mcp_process.poll() is None:
            mcp_process.terminate()
            try:
                mcp_process.wait(timeout=5)
                logger.info("✅ 本地MCP服务已正常关闭")
            except subprocess.TimeoutExpired:
                mcp_process.kill()
                logger.info("⚠️ 本地MCP服务已被强制关闭")


# 程序入口点
if __name__ == "__main__":
    # 设置环境变量，确保使用正确的线程数
    os.environ["OMP_NUM_THREADS"] = str(N_THREADS)
    # 禁用tokenizers并行化以避免警告
    os.environ["TOKENIZERS_PARALLELISM"] = "false"

    # 检查本地MCP服务器文件是否存在
    if os.path.exists(MCP_LOCAL_PATH):
        logger.info(f"检测到本地MCP服务器: {MCP_LOCAL_PATH}")
    elif MCP_ENABLED and not MCP_EXTERNAL and not FLASK_AVAILABLE:
        logger.warning("未检测到本地MCP服务器且Flask不可用，MCP服务将无法启动")

    # 启动主聊天程序
    cuda_chat_main()
